# -*- coding: utf-8 -*-
"""
# --------------------------------------------------------
# @Author : Pan
# @E-mail :
# @Date   : 2025-10-31 10:28:59
# @Brief  : 简单的摄像头播放功能 - 使用gr.Blocks构建
# --------------------------------------------------------
"""
import cv2
import time
import numpy as np
import gradio as gr
import traceback
import copy
import threading
import PIL.Image as Image
from openai import OpenAI
from typing import List, Dict
from pybaseutils import image_utils, file_utils, base64_utils, thread_utils
from pybaseutils.cvutils import video_utils
from pybaseutils.base import list_queue
from app.utils.log import logger
from app.utils import utils

key = "token-abc123"  # API密钥需与启动服务时设置的密钥保持一致
url = "http://192.168.68.102:8000/v1"  # vllm地址
# url = "http://192.168.68.102:40000/v1"  # fastapi chat/completions
# url = "http://0.0.0.0:40000/v1" # fastapi chat/completions
max_turns = 100  # 限制历史记录条数：只保留最近 MAX_TURNS 轮
client = OpenAI(api_key=key, base_url=url)


class WebState:
    code = 1
    msgs = {-1: "等待", 0: "结束", 1: "开始", 2: "继续", 3: "暂停", 4: "播放", 5: "清除"}
    count = 0
    real_time = False


state = WebState()


def finish():
    state.code = 0  # 结束状态
    logger.info(f"state={state.msgs[state.code]}")


def clear_history(message: dict, history=[]):
    state.code = 5  # 清除状态
    logger.info(f"state={state.msgs[state.code]},history= have {len(history)} messages")
    state.code = 0  # 结束状态
    logger.info(f"state={state.msgs[state.code]}")
    return "", "", None


def select_real_time():
    state.real_time = not state.real_time  # 切换实时状态
    logger.info(f"real_time={state.real_time}")
    elem_classes = "button-down" if state.real_time else "button-style"
    return gr.Button(elem_classes=elem_classes)


def display():
    if state.code == 4:
        state.code = 3  # 如何在播放状态，则切换为暂停状态
    else:
        state.code = 4  # 切换为播放状态
    logger.info(f"state={state.msgs[state.code]}")
    return state.msgs[state.code]


def get_history(history=[]):
    """
    获取历史记录中消息和文件
    :param history: 聊天记录，包含用户输入的文本消息和文件，以及系统返回的文本消息和文件
                    history=[{"role": "user",      "content": "用户输入"},
                             {"role": "assistant", "content": "系统结果"}]
    :return: texts文本列表(不含文件), files文件列表
    """
    files, texts = [], []
    for h in history:
        c = h.get("content", [])
        if c and h.get("role", "") == "user" and (isinstance(c, list) or isinstance(c, tuple)):
            files.extend(c)
        else:
            texts.append(h)
    return texts, files


def request_process(question: Dict, frames=[], history=[], reqid=None) -> Dict:
    """
    处理系统返回的文本消息和文件
    level = 0: 表示消息(assistant或者user)，但该消息不参与模型推理
    level = 1: 表示消息(assistant或者user)，且该消息为上下文参与模型推理，默认值
    注意Chatbot只会显示role为assistant和user的消息，其他消息会被忽略
    :param question: 用户输入的文本消息，{"role": "user", "content": "用户输入"}
    :param frames: 用户输入的视频帧列表,[RGB-img]
    :param history: 聊天记录，包含用户输入的文本消息和文件，以及系统返回的文本消息和文件
                    [{"role": "system",      "content": "系统指令"},
                     {"role": "user",        "content": "用户输入"},
                     {"role": "assistant",   "content": "AI回复"},
                     {"role": "runtime",     "content": "运行时信息"},
                     ....
                     ]
    :return: 系统返回的文本消息和文件
    """
    q = copy.deepcopy(question)  # TODO 不要修改message，否则会改变history的内容
    text = q.get("content", "")
    if not text: return {}
    try:
        frames = [Image.fromarray(f) for f in frames]
        prompts = []
        if frames: prompts.append({**q, "content": frames})
        prompts.append({**q, "content": text})
        messages = history + prompts
        messages = utils.gradio2openai_style(messages)
        for i, info in enumerate(messages):
            logger.info("messages[{}]={}".format(i, info))
        messages = base64_utils.serialization(messages, prefix="data:image/jpeg;base64,", use_rgb=True)
        output = client.chat.completions.create(model=None, messages=messages)
        output = output.model_dump()
        result = output['choices'][0]['message']
        # result = {**q, "level": 1, "role": "assistant", "content": [{"type": "text", "text": f"系统收到消息：{text}"}]}
        result = base64_utils.deserialization(result, prefix="data:image/jpeg;base64,", use_rgb=True)
        result = utils.openai2gradio_style(result)
    except Exception as e:
        traceback.print_exc()
        result = {"level": 0, "role": "assistant", "content": f"系统异常: question={q},url={url},error={e}"}
    logger.info("result     ={}".format(result))
    return result


def video_task(producer, consumer, question, his_texts, reqid, winsize, overlap):
    data_list = producer.get_window(winsize=winsize, overlap=overlap, block=True)  # 提取窗口数据
    result = {}
    if data_list:
        logger.info(f"producer qsize={producer.qsize()},thread={threading.get_ident()}")
        times_ = (data_list[0]['time'], data_list[-1]['time'])
        count = [info['count'] for info in data_list]
        frames = [data['frame'] for data in data_list]
        question.update({"time": times_, "count": count})
        result = request_process(question=question, frames=frames, history=his_texts, reqid=reqid)
        result.update({"time": times_, "count": count})
        consumer.put(result)
    return result


def chat_process(message: dict, history=[]):
    """
    处理用户输入的文本消息和文件，返回系统的文本消息和文件
    :param message: 用户输入的消息，包含文本消息和文件
    :param history: 聊天记录，包含用户输入的文本消息和文件，以及系统返回的文本消息和文件
                    history=[{"role": "user",      "content": "用户输入"},
                             {"role": "assistant", "content": "系统结果"}]
    :return: 系统返回的文本和文件消息
    """
    print("-----------" * 10)
    # TODO 限制历史记录条数：只保留最近 max_turns 轮
    if len(history) > max_turns: history = history[-max_turns:]
    reqid = file_utils.get_time(format="id")
    message = {"reqid": reqid, **message}
    logger.info("input message={},history nums={}".format(message, len(history)))
    # TODO 1.用户输入的文本和文件消息
    his_texts, his_files = get_history(history)  # 获得历史记录中的文本消息和文件
    inp_text = message.get("text", "")
    inp_file = message.get("files", [])  # 仅仅支持一个文件(图片或视频)
    reqid = message.get("reqid", None)  # 仅仅支持一个文件(图片或视频)
    if not inp_text: yield "", history, None
    use_file = inp_file[0] if inp_file else his_files[-1] if his_files else ""  # 优先使用用户输入的文件，否则使用历史文件
    if use_file.endswith(".txt"):
        use_file = file_utils.read_file(use_file)
        use_file = file_utils.str2number(use_file)
    logger.info("history files ={}".format(his_files))
    # TODO 2.用户显示的文本和文件消息
    question = {"reqid": reqid, "level": 1, "role": "user", "content": inp_text}
    if inp_file:  # inp_file is list
        history.append({"reqid": reqid, "level": 1, "role": "user", "content": inp_file})
    if inp_text:
        history.append(question)
    # TODO 3.系统处理文本和文件消息
    state.code = 1  # TODO 开始处理
    logger.info(f"state={state.msgs[state.code]},use_file={use_file}")
    result, image = {}, None
    if file_utils.is_image(use_file):
        image = image_utils.read_image(use_file, use_rgb=True)
        result = request_process(question=question, frames=[image], history=his_texts, reqid=reqid)
    elif inp_text and isinstance(use_file, int) or file_utils.is_video(use_file):  # TODO 视频文件或摄像头ID
        max_workers = 2  # 视频处理线程数
        freq = 2  # 视频采样频率，即每秒钟采样帧数
        wintime = 3  # 窗口时间长度，单位秒
        overlap = 0.5  # 窗口重叠比例
        winsize = int(wintime * freq)  # 窗口帧数
        # realtime = True  # 是否实时处理视频帧，实时处理会丢弃很多视频帧，非实时会处理所有视频帧，但视频播放会很卡顿
        # realtime = isinstance(use_file, int) or realtime  # 使用摄像头，建议实时处理视频帧
        producer = list_queue.Queue(maxsize=winsize)  # 生产者
        consumer = list_queue.Queue(maxsize=winsize)  # 消费者
        # TODO 打开视频文件或摄像头
        w, h, num_frames, fps = image_utils.get_video_info(use_file)
        interval = int(fps / freq) if fps > 0 else 1
        video_cap = video_utils.video_iterator(use_file, save_video=None)
        pool = thread_utils.ThreadPool(max_workers=max_workers)
        future = None
        # TODO 主线程处理视频帧
        for data_info in video_cap:
            while state.code == 3:  # 暂停状态，等待继续状态
                time.sleep(0.05)
            if state.code == 0: break
            state.code = 4  # 播放状态
            image = data_info["frame"][:, :, ::-1]  # BGR to RGB
            image = image_utils.resize_image(image, size=(None, 640))
            data_info["frame"] = image
            finish = data_info["finish"]
            # TODO 视频抽帧，放入生产者队列，非实时处理会阻塞，实时处理会丢弃旧的视频帧
            if data_info['count'] % interval == 0:
                producer.put(data_info, block=not state.real_time)
            # TODO 启动视频处理线程
            if producer.qsize() >= winsize:
                winsize = producer.qsize() if finish else winsize
                future = pool.submit(video_task, *(producer, consumer, question, his_texts, reqid, winsize, overlap))
            # time.sleep(0.01)
            # TODO 从消费者队列中获取系统分析结果(注意：返回结果可能是无序，建议根据发送时间进行排序)
            while consumer.qsize() > 0 or (finish and future and not future.done()):
                result = consumer.pop(block=False, timeout=0.005)
                if state.code == 0: break
                if result:
                    times_ = result['time']
                    info = [{"reqid": reqid, "level": 0, "role": "assistant", "content": f"- Time:{times_}s"}, result]
                    history += info
                    # his_texts.append(question) # TODO 是否添加历史数据,可能效果变差
                    # his_texts.append(result)
                    yield "", history, image  # RGB image
            result = {}  # 避免最后结果重复提交
            yield "", history, image  # RGB image
    else:
        result = request_process(question=question, history=his_texts, reqid=reqid)
    if result:
        history.append(result)
    state.code = 0  # TODO 处理完成
    logger.info(f"state={state.msgs[state.code]},use_file={use_file}")
    yield "", history, image


CSS = """
.web_title {
    text-align: center;
    margin: 5px 0;                /* 上下间距5px,左右0px */
    padding: 0px;                 /* 内边距 */
    color: #000000;
}
.web_footer {                     /* 自定义页脚 */
    text-align: center;           /*           */
    margin-top: 100px;
    color: #666;
}
.gradio-container {
    width: 100% !important;
    max-width: none !important;
}
.gradio-image {
    height: 480px;
}
.gradio-image img {
    height: 480px;
    object-fit: contain !important;
}
.row-align {                 
    display: flex;
    justify-content: center;
    align-items: center;
}
.button-style {
    text-align: center;
    font-size: 13px;
    height: 30px;
    max-width: 100px;
    min-width: unset !important;   /* 取消最小宽度限制 */
    border-radius: 20px;
}
.button-down {
    text-align: center;
    font-size: 13px;
    height: 30px;
    max-width: 100px;
    min-width: unset !important;   /* 取消最小宽度限制 */
    border-radius: 20px;
    background: #FF8C00 !important;
}
"""


def ui_chatbot(name=""):
    with gr.Tab(label=name):
        with gr.Row():
            out = gr.Chatbot(label="聊天记录",
                             height=600,
                             scale=1,
                             show_copy_button=True,
                             render_markdown=True,
                             sanitize_html=False,
                             avatar_images=["data/assets/user.png", "data/assets/system.png"],
                             type="messages",
                             )
            vid = gr.Image(label="视频", height=600, scale=1, streaming=True, elem_classes=["gradio-image"])
        inp = gr.MultimodalTextbox(interactive=True,
                                   placeholder="输入文本消息，或者上传图片/视频文件",
                                   show_label=False,
                                   file_types=["image", "video", "audio", ".pdf", ".txt"],
                                   )
        with gr.Row(elem_classes="row-align"):
            stp = gr.Button("停止", elem_classes="button-style", size="md")
            clc = gr.Button("清除", elem_classes="button-style", size="md")
            rtb = gr.Button("实时", elem_classes="button-style", size="md")
        inp.submit(fn=chat_process, inputs=[inp, out], outputs=[inp, out, vid], show_progress_on=[out])
        stp.click(fn=finish)
        clc.click(fn=clear_history, inputs=[inp, out], outputs=[inp, out, vid])
        rtb.click(fn=select_real_time, outputs=[rtb])


def ui_titles(name=""):
    # 自定义标题
    ui = gr.HTML(f""" 
                <h1 class="web_title">{name}</h1>
                """)
    return ui


def ui_footer():
    # 自定义页脚
    ui = gr.HTML("""<div class="web_footer">
                    Created by PanJinquan | Contact: pan_jinquan@163.com
                    </div>
                """)
    return ui


def create_app(title="AI小工具"):
    """
    :param title:
    :return:
    """
    with gr.Blocks(title=title, css=CSS) as app:
        ui_titles(name=title)
        ui_chatbot(name="聊天")
        ui_footer()
    return app


if __name__ == "__main__":
    # 创建并启动界面
    app = create_app()
    app.launch(server_name="0.0.0.0", share=False, show_error=True)
